package com.wang.madness;

import com.alibaba.fastjson.JSON;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import org.assertj.core.util.Lists;

public class TreeUtils {

  static
  class Node {

    /**
     * id
     */
    private int id;
    /**
     * 父级id
     */
    private int pid;
    /**
     * 节点名称
     */
    private String name;
    /**
     * 子节点
     */
    private List<Node> sub = new ArrayList<>();

    public Node(int id, int pid) {
      this.id = id;
      this.pid = pid;
    }

    public Node(int id, int pid, String name) {
      this(id, pid);
      this.name = name;
    }

    public int getId() {
      return id;
    }

    public void setId(int id) {
      this.id = id;
    }

    public int getPid() {
      return pid;
    }

    public void setPid(int pid) {
      this.pid = pid;
    }

    public String getName() {
      return name;
    }

    public void setName(String name) {
      this.name = name;
    }

    public List<Node> getSub() {
      return sub;
    }

    public void setSub(List<Node> sub) {
      this.sub = sub;
    }
  }

  public static List<Node> buildTree(List<Node> nodes) {
    // 按父级id分组
    Map<Integer, List<Node>> sub = nodes.stream().filter(node -> node.getPid() != 0).collect(Collectors.groupingBy(Node::getPid));
    // 给每个node,根据上面的分组设置children
    nodes.forEach(node -> node.setSub(sub.get(node.getId())));
    // 只保留根节点
    return nodes.stream().filter(node -> node.getPid() == 0).collect(Collectors.toList());
  }

  public static void main(String[] args) {
    Node dennis = new Node(1, 0, "dennis");
    Node calm = new Node(2, 0, "calm");
    Node daughter = new Node(3, 1, "daughter");
    Node grandson = new Node(4, 3, "grandson");
    Node son = new Node(5, 2, "son");
    List<Node> nodes = Lists.newArrayList(dennis, calm, daughter, son, grandson);
    List<Node> tree = buildTree(nodes);
    System.out.println(JSON.toJSONString(tree));
  }

}
